{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TwCbR5P4c3vA"
   },
   "source": [
    "<img src='http://www-scf.usc.edu/~ghasemig/images/sharif.png' alt=\"SUT logo\" width=200 height=200 align=left class=\"saturate\" >\n",
    "\n",
    "<br>\n",
    "<font face=\"Times New Roman\">\n",
    "<div dir=ltr align=center>\n",
    "<font color=0F5298 size=7>\n",
    "    ML for Bioinformatics <br>\n",
    "<font color=2565AE size=5>\n",
    "    Computer Engineering Department <br>\n",
    "<font color=3C99D size=5>\n",
    "    Homework 3: Practical - Variational Autoencoder (VAE) <br>\n",
    "<font color=696880 size=4>\n",
    "    Mahdi Manouchehry (mahdimanouchehry14@gmail.com) <br>\n",
    "    \n",
    "____\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Full Name : \n",
    "### Student Number : \n",
    "__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GqG8IST7m-_F"
   },
   "source": [
    "In this [link](https://lilianweng.github.io/posts/2018-08-12-vae/), you can find out more about autoencoders in general and a good explanation about variational autoencoders as well, known as VAEs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "REwELF9ikTgb"
   },
   "outputs": [],
   "source": [
    "import torch \n",
    "import torch.nn as nn \n",
    "import torch.nn.functional as F \n",
    "from torchvision import datasets, transforms \n",
    "from torch.utils.data import DataLoader \n",
    "import matplotlib.pyplot as plt \n",
    "\n",
    "\n",
    "#Add Additional libraries here\n",
    "\n",
    "\n",
    "torch.manual_seed(42)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gK5o3WmkdQp8"
   },
   "source": [
    "The MNIST dataset is a large database of handwritten digits that is commonly used for training and testing various image processing and machine learning models. It contains 60,000 training images and 10,000 testing images of size (28, 28) and their corresponding labels from 0 to 9."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "a3AJpXKwBECM"
   },
   "outputs": [],
   "source": [
    "# Load the MNIST dataset\n",
    "data = datasets.MNIST('data', train=True, download=True,\n",
    "                   transform=transforms.ToTensor())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GOtIKzeXBG5C"
   },
   "source": [
    "# Data Visualization\n",
    "\n",
    "Show 3 random samples from the dataset.\n",
    "\n",
    "*hint: you can use \"cv2\" library for reading images but any other method is acceptable.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1CvuX1ivBKu_"
   },
   "outputs": [],
   "source": [
    "#################### Problem 01  ####################\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xyRmhulfTAa0"
   },
   "source": [
    "This code defines a VAE class for a variational autoencoder (VAE) model, a generative model with an encoder-decoder based structure. The encoder compresses an input to a latent vector, which is drawn from a distribution learned by the encoder outputs. The decoder reconstructs the input from the latent vector, and the model optimizes the reconstruction quality and the difference between the latent distribution and a prior distribution. A VAE class usually specifies the encoder and decoder networks, the sampling layer, and the loss function.\n",
    "\n",
    "# Complete the following class to define the VAE structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "_IU5StXrksLp"
   },
   "outputs": [
    {
     "ename": "SyntaxError",
     "evalue": "invalid syntax (<ipython-input-2-d030e189c0d9>, line 9)",
     "output_type": "error",
     "traceback": [
      "\u001b[1;36m  File \u001b[1;32m\"<ipython-input-2-d030e189c0d9>\"\u001b[1;36m, line \u001b[1;32m9\u001b[0m\n\u001b[1;33m    self.hidden_dim = pass\u001b[0m\n\u001b[1;37m                      ^\u001b[0m\n\u001b[1;31mSyntaxError\u001b[0m\u001b[1;31m:\u001b[0m invalid syntax\n"
     ]
    }
   ],
   "source": [
    "# Define the VAE architecture\n",
    "class VAE(nn.Module):\n",
    "    \n",
    "    def __init__(self, hidden_size, latent_size):\n",
    "        super(VAE, self).__init__()\n",
    "        #################### Problem 02 - 03 ####################\n",
    "        \n",
    "        # Assign the dimensions of the hidden layer and the latent space\n",
    "        self.hidden_dim = pass\n",
    "        self.latent_dim = pass\n",
    "    \n",
    "        # Deifne Encoder layers: map the input data to the mean and log variance of the latent space\n",
    "        # Your Code\n",
    "        \n",
    "        # Define Decoder layers: map the latent space to the reconstructed output data\n",
    "        # Your Code\n",
    "        \n",
    "        \n",
    "\n",
    "    def encode(self, x):\n",
    "        #################### Problem 04  ####################\n",
    "        '''\n",
    "        Encode the input to the latent space using linear transformation layers and a ReLU activation function.\n",
    "        '''\n",
    "        \n",
    "        # Your Code\n",
    "    \n",
    "        return # Return the mean and log variance of the latent space\n",
    "    \n",
    "    \n",
    "    \n",
    "    def reparameterize(self, mu, logvar):\n",
    "        #################### Problem 05  ####################\n",
    "        '''\n",
    "        # 1. Compute the standard deviation from the log variance\n",
    "        # 2. Sample a random tensor from the standard normal distribution with the same shape as std\n",
    "        # 3. return a sample from the latent space using the reparameterization trick\n",
    "        '''\n",
    "        \n",
    "        #Your Code\n",
    "        \n",
    "        return  # Return a sample from the latent space using the reparameterization trick\n",
    "    \n",
    "    \n",
    "    \n",
    "    def decode(self, z):\n",
    "        #################### Problem 06  ####################\n",
    "        '''\n",
    "         Decode the latent space to the output space through:\n",
    "         1. Apply a linear transformation and a ReLU activation to the latent space\n",
    "         2. Apply a linear transformation and a sigmoid activation to the hidden layer\n",
    "        '''\n",
    "    \n",
    "        # Your Code\n",
    "           \n",
    "        return  # Return result\n",
    "    \n",
    "    \n",
    "    \n",
    "    def forward(self, x):\n",
    "        #################### Problem 07  ####################\n",
    "        '''\n",
    "        The forward propagation of your model:\n",
    "        1. Encode the input data and get the mean and log variance of the latent space\n",
    "        2. Reparameterize the mean and log variance and get a sample from the latent space\n",
    "        3. Return the decoded samples, mu and logvar\n",
    "        '''\n",
    "        \n",
    "        # Your Code\n",
    "        \n",
    "        return # return results\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1I7MqCoHZ-r2"
   },
   "source": [
    "# Define the Loss Function\n",
    "Compute the binary cross entropy between the reconstructed data and the original data.\n",
    "\n",
    "Also, you need to compute the KL divergence between the distribution of the latent space and a standard normal distribution.\n",
    "\n",
    "The total loss would be the sum of these two losses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Tc42rsPQn7ZN"
   },
   "outputs": [],
   "source": [
    "def loss_function(recon_x, x, mu, logvar):\n",
    "    #################### Problem 08  ####################\n",
    "    '''\n",
    "    Compute the binary cross entropy between the reconstructed data and the original data\n",
    "    Compute the KL divergence between the Gaussian distribution of the latent space and a standard normal distribution\n",
    "    '''\n",
    "    \n",
    "    # Your code\n",
    "    \n",
    "    return # Return the sum of binary cross entropy and KL divergence as the total loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NEh19Yy-aTup"
   },
   "source": [
    "# Training the Model\n",
    "\n",
    "Here we have provided a function which trains the model using the specified input arguments. (Make sure to understand the code!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "M3o6xuOFoQJk"
   },
   "outputs": [],
   "source": [
    "def train(model, optimizer, train_loader, device):\n",
    "    model.train()\n",
    "    train_loss = 0\n",
    "    for batch_idx, (data, _) in enumerate(train_loader):\n",
    "        data = data.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        recon_batch, mu, logvar = model(data)\n",
    "        loss = loss_function(recon_batch, data, mu, logvar)\n",
    "        loss.backward()\n",
    "        train_loss += loss.item()\n",
    "        optimizer.step()\n",
    "    print('====> Epoch: {} Average loss: {:.4f}'.format(epoch, train_loss / len(train_loader.dataset)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2czyDhzyoq5S"
   },
   "outputs": [],
   "source": [
    "#################### Problem 09  ####################\n",
    "  # Set the batch size for loading the data\n",
    "####################################################\n",
    "batch_size = pass\n",
    "\n",
    "\n",
    "#################### Problem 10  ####################\n",
    "  #Create a data loader object that shuffles and batches the data\n",
    "####################################################\n",
    "train_loader = pass\n",
    "\n",
    "\n",
    "#Choose the device (CPU or GPU) based on availability\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "#Create an instance of the VAE model and move it to the device\n",
    "model = VAE().to(device)\n",
    "\n",
    "\n",
    "#################### Problem 11  ####################\n",
    "#Create an instance of the Adam optimizer with a learning rate(try different learning rate)\n",
    "#Set the number of epochs to train the model \n",
    "####################################################\n",
    "optimizer = pass\n",
    "epochs = pass\n",
    "\n",
    "for epoch in range(1, epochs+1):\n",
    "    train(model, optimizer, train_loader, device)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9CsxSbLqilxq"
   },
   "source": [
    "# Generative Model\n",
    "Now that you have trained your model, you can generate new images by sampling from the latent space and decoding them as below: (make sure to understand how the code works!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VcPha9JZij2k"
   },
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    z = torch.randn(batch_size, 20).to(device)\n",
    "    sample = model.decode(z).cpu()\n",
    "    \n",
    "    # Plot the generated images\n",
    "    fig, ax = plt.subplots(nrows=5, ncols=10, figsize=(10,5))\n",
    "    for i in range(5):\n",
    "        for j in range(10):\n",
    "            ax[i][j].imshow(sample[i*10+j].reshape(28, 28), cmap='gray')\n",
    "            ax[i][j].axis('off')\n",
    "    plt.show()"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
